/* Copyright 2012 Tobias Marschall
 * 
 * This file is part of CLEVER.
 * 
 * CLEVER is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * CLEVER is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <limits>
// #include <stdio.h>
// #include <io.h>
#include <cassert>
#include <iomanip>

#include <boost/program_options.hpp>
#include <bamtools/api/BamReader.h>

#include "InsertLengthDistribution.h"
#include "GroupWiseBamReader.h"

using namespace std;
namespace po = boost::program_options;

int phred_base = 33;

long long skipped_inside_clipping = 0;

void usage(const char* name, const po::options_description& options_desc) {
	cerr << "Usage: " << name << " [options] <insert-length-dist>" << endl;
	cerr << endl;
	cerr << "<insert-length-dist>: Insert length distribution as obtained by running" << endl;
	cerr << "                      insert-length-histogram on a BAM file." << endl;
	cerr << endl;
	cerr << "Reads BAM format from stdin and outputs a list of plausible alignments" << endl;
	cerr << "(one per line) of read pairs to stdout in the following format:" << endl;
	cerr << endl;
	cerr << "<read-name> <pair-nr> <phred-sum1> <chromosome1> <start1> <end1> <strand1> <phred-sum2> <chromosome2> <start2> <end2> <strand2> <aln-pair-prob> <aln-pair-prob-inslength>" << endl;
	cerr << endl;
	cerr << "Where <aln-pair-prob> is the probability that this alignment pair is " << endl;
	cerr << "correct estimated based on the alignment quality while the probabilities " << endl;
	cerr << "given in <aln-pair-prob-inslength> are estimates based on alignment quality AND" << endl;
	cerr << "insert length." << endl;
	cerr << endl;
	cerr << "<startX> and <endX> coordinates are 1-based and inclusive (i.e. closed intervals)" << endl;
	cerr << "and give the region on the reference the respective read was aligned to." << endl;
	cerr << endl;
	cerr << "The probability estimates are based on the alignment score (AS) tags in the " << endl;
	cerr << "BAM input. Therefore, these tags must be present. BAM files without these" << endl;
	cerr << "tags should be piped through add-score-tags-to-bam to compute them." << endl;
	cerr << endl;
	cerr << "Unmapped reads are ignored. Alignment pairs for which both alignments have the same " << endl;
	cerr << "orientation (strandedness) or map to different chromosomes are not reported. " << endl;
	cerr << endl;
	cerr << "IMPORTANT: Assumes that, in the BAM input, alignments belonging to the same read" << endl;
	cerr << "           are grouped, i.e. are given in subsequent lines." << endl;
	cerr << endl;
	cerr << "NOTE: When using BWA, multiple alignment locations are by default stored in XA tags" << endl;
	cerr << "      in a single line (i.e. one line per read rather than per alignment). Can be " << endl;
	cerr << "      converted using the script xa2multi.pl that comes with BWA." << endl;
	cerr << endl;
	cerr << options_desc << endl;
	exit(1);
}

/** Attaches probabilities corresponding with the 'AS' scores of the alignments. */
vector<double>* compute_posteriors(const vector<BamTools::BamAlignment>& alignments) {
	double probsum = 0.0;
	vector<double>* probs = new vector<double>();
	typedef vector<BamTools::BamAlignment>::const_iterator aln_iter_t;
	for (aln_iter_t it=alignments.begin(); it!=alignments.end(); ++it) {
		uint32_t as = 0;
		if (it->GetTag("AS", as)) {
			double p = pow(10.0,as/-10.0);
			probs->push_back(p);
			probsum += p;
		} else {
			assert(false);
		}
	}
	if (probsum > 0.0) {
		for (size_t i=0; i<probs->size(); ++i) {
			probs->at(i) /= probsum;
		}
	} else {
		for (size_t i=0; i<probs->size(); ++i) {
			probs->at(i) = 1.0/alignments.size();
		}
	}
	return probs;
}

int phred_sum(const BamTools::BamAlignment& aln) {
	const string& phred = aln.Qualities;
	int result = 0;
	for (size_t i=0; i<phred.size(); ++i) {
		result += phred[i] - phred_base;
	}
	return result;
}

typedef struct aln_pair_t {
	const BamTools::BamAlignment* aln1;
	const BamTools::BamAlignment* aln2;
	double pair_prob;
	double pair_prob_ins_length;
	aln_pair_t(const BamTools::BamAlignment* aln1, const BamTools::BamAlignment* aln2, double pair_prob, double pair_prob_ins_length) : aln1(aln1), aln2(aln2), pair_prob(pair_prob), pair_prob_ins_length(pair_prob_ins_length) {}
} aln_pair_t;

/** Given an output file (a plain file object) and a list of alignments belonging
to the same read pair, outputs one line per alignment pair in the format:
"<read-name> <chromosome1> <pos1> <strand1> <chromosome2> <pos2> <strand2> <aln-pair-prob> <aln-pair-prob-norm>".*/
void process_read_allpairs(const BamTools::RefVector& ref_vector, const vector<BamTools::BamAlignment>& alignments1, const vector<BamTools::BamAlignment>& alignments2, const InsertLengthDistribution* ild) {
	vector<double>* posteriors1 = compute_posteriors(alignments1);
	vector<double>* posteriors2 = compute_posteriors(alignments2);
	assert(posteriors1->size() == alignments1.size());
	assert(posteriors2->size() == alignments2.size());
	// contains an entry for every pair that should be printed
	vector<aln_pair_t> all_pairs;
	double p_sum = 0.0;
	double p_ins_sum = 0.0;
	for (size_t i=0; i<alignments1.size(); ++i) {
		const BamTools::BamAlignment& read_aln1 = alignments1[i];
		assert(read_aln1.CigarData.size() > 0);
		double posterior1 = posteriors1->at(i);
		assert(read_aln1.IsMapped());
		for (size_t j=0; j<alignments2.size(); ++j) {
			const BamTools::BamAlignment& read_aln2 = alignments2[j];
			assert(read_aln2.CigarData.size() > 0);
			double posterior2 = posteriors2->at(j);
			assert(read_aln2.IsMapped());
			// skip if strandedness is the same
			if (read_aln1.IsReverseStrand() == read_aln2.IsReverseStrand()) continue;
			// skip if both reads map to different chromosomes
			if (read_aln1.RefID != read_aln2.RefID) continue;
			double p = posterior1*posterior2;
			double p_ins = p;
			const BamTools::BamAlignment& left = (read_aln1.Position<=read_aln2.Position)?read_aln1:read_aln2;
			const BamTools::BamAlignment& right = (read_aln1.Position<=read_aln2.Position)?read_aln2:read_aln1;
			if (((left.CigarData[left.CigarData.size()-1].Type == 'S') && (left.CigarData[left.CigarData.size()-1].Length > 2)) ||
				((right.CigarData[0].Type == 'S') && (right.CigarData[0].Length > 2))) {
				skipped_inside_clipping += 1;
				delete posteriors1;
				delete posteriors2;
				return;
			}
			if (ild != 0) {
				int insert_length = right.Position - left.GetEndPosition();
				p_ins *= ild->probability(insert_length);
			}
			p_sum += p;
			p_ins_sum += p_ins;
			all_pairs.push_back(aln_pair_t(&left, &right, p, p_ins));
		}
	}
	if ((p_sum > 0.0) && (p_ins_sum > 0.0)) {
		int n = 0;
		for (size_t i=0; i<all_pairs.size(); ++i) {
			aln_pair_t& p = all_pairs[i];
			cout << p.aln1->Name << " " << (n++) << " "
			     << phred_sum(*(p.aln1)) << " " << ref_vector[p.aln1->RefID].RefName << " " << p.aln1->Position+1 << " " << p.aln1->GetEndPosition() << " " << (p.aln1->IsReverseStrand()?'-':'+') << " "
			     << phred_sum(*(p.aln2)) << " " << ref_vector[p.aln2->RefID].RefName << " " << p.aln2->Position+1 << " " << p.aln2->GetEndPosition() << " " << (p.aln2->IsReverseStrand()?'-':'+') << " "
		         << setprecision(16) << p.pair_prob/p_sum << " " << p.pair_prob_ins_length/p_ins_sum << endl;
		}
	}
	delete posteriors1;
	delete posteriors2;
}

void check_as_tags(const vector<BamTools::BamAlignment>& alignments) {
	vector<BamTools::BamAlignment>::const_iterator it = alignments.begin();
	for (; it != alignments.end(); ++it) {
 		uint32_t as_tag = 0;
 		if (!it->GetTag("AS", as_tag)) {
 			ostringstream oss;
 			oss << "Error: No AS tag present for read \"" << it->Name << "\"." << endl;
 			throw std::runtime_error(oss.str());
 		}
	}
}

int main(int argc, char* argv[]) {
	// PARAMETERS
	bool skip_non_xa = false;
	bool ignore_wrong_x_tags = false;

	po::options_description options_desc("Allowed options");
	options_desc.add_options()
		("skip_non_xa,X", po::value<bool>(&skip_non_xa)->zero_tokens(), "Skip reads for which other alignments exist (i.e. X0+X1>1), but no XA tag is present. Turn on when using BWA.")
		("ignore_wrong_X_tags,I", po::value<bool>(&ignore_wrong_x_tags)->zero_tokens(), "Do not abort when wrongly set X0/X1 tags are encountered.")
		("phred_base,p", po::value<int>(&phred_base)->default_value(33), "Value to substract from ASCII code to get the PHRED quality.")
	;
	
	if (isatty(fileno(stdin)) || (argc<2)) {
		usage(argv[0], options_desc);
	}
	string insert_length_histogram(argv[argc-1]);
	argc -= 1;

	po::variables_map options;
	try {
		po::store(po::parse_command_line(argc, argv, options_desc), options);
		po::notify(options);
	} catch(exception& e) {
        cerr << "error: " << e.what() << "\n";
        return 1;
    }

	GroupWiseBamReader* bam_reader = 0;
	InsertLengthDistribution* insert_length_distribution = 0;
	// count the number of alignments skipped because of missing XA tags
	long long skipped_by_xa = 0;
	long long bad_x_tags = 0;
	try {
		bam_reader = new GroupWiseBamReader("/dev/stdin", true);
		bam_reader->enableProgressMessages(cerr, 200000);
		insert_length_distribution = new InsertLengthDistribution(insert_length_histogram);
		while ( bam_reader->hasNext() ) {
			bam_reader->advance();
			if (bam_reader->isFirstUnmapped() || bam_reader->isSecondUnmapped()) continue;
			if (skip_non_xa && bam_reader->hasMultipleMappingsButNoXA()) {
				skipped_by_xa += 1;
				continue;
			}
			if (bam_reader->hasInconsistentXTags()) {
				if (ignore_wrong_x_tags) {
					bad_x_tags += 1;
				} else {
					cerr << bam_reader->getInconsistentXTagsMessage() << endl;
					cerr << "  This can have several reasons, for example:" << endl;
					cerr << "  1) You use BWA and have forgotten to run xa2multi.pl (distributed with BWA)." << endl;
					cerr << "  2) The used read aligner does not write X0 and X1 tags correctly." << endl;
					if (!skip_non_xa) {
						cerr << "  3) If you use BWA alignments, remember to use option -X." << endl;
					}
					cerr << endl;
					cerr << "To ignore this inconsistency, use option -I." << endl;
					return 1;
				}
			}
			const vector<BamTools::BamAlignment>& alignments1 = bam_reader->getAlignmentsFirst();
			const vector<BamTools::BamAlignment>& alignments2 = bam_reader->getAlignmentsSecond();
			check_as_tags(alignments1);
			check_as_tags(alignments2);
			process_read_allpairs(bam_reader->getReferenceData(), alignments1, alignments2, insert_length_distribution);
		}
	} catch(exception& e) {
        cerr << "Error: " << e.what() << "\n";
        return 1;
    }
	if (bad_x_tags > 0) {
		cerr << "Skipped " << bad_x_tags << " reads with inconsistent X0/X1 tags (as requested by option -I)." << endl;
	}
	if (skipped_by_xa > 0) {
		cerr << "Skipped " << skipped_by_xa << " ambiguously mapped reads for which no XA tag was present (as requested by option -X)." << endl;
	}
	if (bam_reader->getSkippedDuplicates() > 0) {
		cerr << "Skipped " << bam_reader->getSkippedDuplicates() << " duplicate alignments." << endl;
	}
	if (skipped_inside_clipping > 0) {
		cerr << "Skipped " << skipped_inside_clipping << " reads with alignments soft-clipped on the inside of an alignment pair!" << endl;
	}
	if (insert_length_distribution!=0) {
		delete insert_length_distribution;
	}
	return 0;
}
